#include "rsa_32b.h"
#include <stdlib.h>
#include <time.h>
#include <string.h>
#include <stdio.h>
#include "base64.h"
#include "config_parse/config_parse.h"
#include "stringlib.h"
#include "sha/sha1.h"

#define BAD_WAY 1
#define VERSION_RSA_32B "v0"

static int gcd(int m, int n) {
    int r;
    if (m < n) {
        m^=n;
        n^=m;
        m^=n;
    }
    while ((m%n) != 0) {
        r = m%n;
        m = n;
        n = r; 
    }
    return n;
}

static int eIsStandard(int e,int fai_n) {
    if(e>1 && e<fai_n) {
        if (gcd(e, fai_n) == 1) return 1;
        else return 0;
    }
    return 0;
}

// 计算模逆元
static int mod_inverse(int a, int m) {
    int m0 = m, t, q;
    int x0 = 0, x1 = 1;

    if (m == 1) return 0;
    while (a > 1) {
        q = a / m;
        t = m;

        m = a % m;
        a = t;
        t = x0;

        x0 = x1 - q * x0;
        x1 = t;
    }
    if (x1 < 0) x1 += m0;

    return x1;
}

static int countD(int e, int fai_n) {
    long long k=0, d;
    for (k = 0; (k*fai_n + 1) % e != 0; k++);
    // d * e ≡ 1 (mod Φ(n))
    d = (k*fai_n + 1) / e;
    // printf("k = %d\n",k);
    return d;
}

// 模幂运算
static int mod_exp(long long base, int exp, int mod) {
    if (base == 0) return mod;
    if (base == mod) return 0;
    // printf("base: %ld, ", base);
    long long result = 1;
    while (exp > 0) {
        if (exp % 2 == 1) result = (result * base) % mod;
        base = (base * base) % mod;
        exp /= 2;
    }
    // printf("result: %ld\n", result);
    return result;
}

static int isLittleEndian() {
    union un
    {
        char c;
        int i;
    } u = {1};
    return u.c;
}

static int isPrime(int p) {
    if (p<=2) return 0;
    int i;
    for(i=2; i<p; i++) {
        if(p%i==0) break;
    }
    if (i<p) return 0;  // not prime
    else return 1;      // is prime
}

// struct rsa32b_key
// {
//     int q;     // big prime number
//     int p;     // big prime number
//     int n;     // p * q
//     // int fai_n; // φ(n) = (p-1)*(q-1)
//     int e;     // 1 < e < φ(n), e is Prime with φ(n), normal will set 65537
//     int d;     // d·e≡1 mod φ(n)
//     // expand
//     int padding;    
// };

enum rsa_padding {
    RSA_NO_PADDING = 0,
    RSA_PKCS1_PADDING,
    RSA_PKCS1_OAEP_PADDING,
};

int rsa32b_key_create(struct rsa32b_key *rkey, int p, int q, int e) {
    srand(time(NULL));
    if (p) {
        rkey->p = p;
    } else {
        rkey->p = rand();
        // if (rkey->p <= 0x7fff) rkey->p*=rkey->p;
        while (rkey->p > 2) {
            if (isPrime(rkey->p)) break;
            rkey->p--;
        }
    }
    if (q) {
        rkey->q = q;
    } else {
        rkey->q = rand();
        // if (rkey->q <= 0x7fff) rkey->q*=rkey->q;
        while (rkey->q > 2) {
            if (isPrime(rkey->q)) break;
            rkey->q--;
        }
    }
    // printf("p = %d\n", rkey->p);
    // printf("q = %d\n", rkey->q);
    if (!isPrime(rkey->p)) return -1;
    if (!isPrime(rkey->q)) return -2;

    rkey->n = rkey->p * rkey->q;
    int fai_n = (rkey->p-1) * (rkey->q-1);
    if ((rkey->n/rkey->p!=rkey->q) || (fai_n/(rkey->p-1)!=(rkey->q-1))) return -3;
    
    if (e) rkey->e = e;
    else rkey->e = 65537;
    if(!eIsStandard(rkey->e, fai_n)) return -4;
    // rkey->d = mod_inverse(rkey->e, fai_n);
    rkey->d = countD(rkey->e, fai_n);
    rkey->padding = RSA_NO_PADDING;
    printf("p = %d, q = %d\n", rkey->p, rkey->q);
    printf("n = %d, fai_n = %d\n", rkey->n, fai_n);
    printf("e = %d, d = %d\n", rkey->e, rkey->d);
    return 0;
}

// RSA 加密
int rsa_encrypt_simple(int plaintext, int e, int n) {
    return mod_exp(plaintext, e, n);
}

// RSA 解密
int rsa_decrypt_simple(int ciphertext, int d, int n) {
    return mod_exp(ciphertext, d, n);
}

static int rsa_no_padding(char *instr, int len, int align_bits) {
    int pl = 8 - len%(align_bits/8);
    if (!pl) return pl;
    for (instr+=len;pl>0;pl--,instr++,len++) {
        *instr = 0;
    }
    return len;
}

static int rsa_remove_no_padding(char *instr, int len, int align_bits) {
    while (instr[len-1] == 0) len--;
    return len;
}

/* Swap bytes in 32 bit value. */
#define __bswap_32(x)    \
     ((((x) & 0xff000000) >> 24)  \
     | (((x) & 0x00ff0000) >>  8) \
     | (((x) & 0x0000ff00) <<  8) \
     | (((x) & 0x000000ff) << 24))

static int rsa_encrypt(char * pla, char* out, int len, int e, int n) {
    int *pin = (int *)malloc(len+4);
    char *mc_out = malloc(len+len/4+5); // (len+4)*5/4
    int *pout = (int *)mc_out;
    if (!pin || !mc_out) {
        free(pin);
        free(mc_out);
        return -1;
    }
    memcpy(pin, pla, len);
    len = rsa_no_padding((char*)pin, len, 32);
    int plen = len/sizeof(int);
    for (int i=0; i<plen; i++,len++) {
    #if BAD_WAY
        if (pin[i]/n >255) {
            free(pin);
            free(mc_out);
            return -1;   // please use larger q and p
        }
        mc_out[i*5] = pin[i]/n;
        pout = (int *)(mc_out+i*5+1);
        pout[0] = mod_exp(pin[i], e, n);
    #else 
        if (pin[i] > n) {
            free(pin);
            free(mc_out);
            return -1;
        }
        pout[i] = mod_exp(pin[i], e, n);
    #endif
    }
    // (n+2)*4/3
    base64_encode(mc_out, len, out, &len);

    free(pin);
    free(mc_out);
    return len;
}

static int rsa_decrypt(char *cip, char* out, int len, int d, int n) {
    int *pin = (int *)cip;
    int *pout = (int *)out;
    char *mc_in = (char*)malloc(len);
    if (!mc_in || base64_decode(cip, len, mc_in, &len)) {
        free(mc_in);
        return -1;
    }
    for (int i=0; i<len/(sizeof(int)+1); i++) {
    #if BAD_WAY
        pin = (int *)(mc_in+i*5+1);
        pout[i] = mod_exp(pin[0], d, n) + mc_in[i*5]*n;
    #else
        pout[i] = mod_exp(pin[i], d, n);
    #endif
    }
    free(mc_in);
    return rsa_remove_no_padding(out, len, 32);
}

int rsa32b_public_encrypt(struct rsa32b_key *rkey, char *pla, char *out, int len) {
    if (!rkey || !rkey->e || !rkey->n || !pla || !out || !len) return -1;
    return rsa_encrypt(pla, out, len, rkey->e, rkey->n);
}

int rsa32b_public_decrypt(struct rsa32b_key *rkey, char *cip, char *out, int len) {
    if (!rkey || !rkey->e || !rkey->n || !cip || !out || !len) return -1;
    return rsa_decrypt(cip, out, len, rkey->e, rkey->n);
}

int rsa32b_private_encrypt(struct rsa32b_key *rkey, char *pla, char *out, int len) {
    if (!rkey || !rkey->d || !rkey->n || !pla || !out || !len) return -1;
    return rsa_encrypt(pla, out, len, rkey->d, rkey->n);
}

int rsa32b_private_decrypt(struct rsa32b_key *rkey, char *cip, char *out, int len) {
    if (!rkey || !rkey->d || !rkey->n || !cip || !out || !len) return -1;
    return rsa_decrypt(cip, out, len, rkey->d, rkey->n);
}
/*
// out need space: (((len+4)*5/4)+2)*4/3
int test_encrypt(char * pla, char* out, int len, int e, int n) {
    printf("%d line: %d\n", __LINE__, len);
    int *pin = (int *)malloc(len+4);
    char *mc_out = malloc(len+len/4+5); // (len+4)*5/4
    int *pout = (int *)mc_out;
    memcpy(pin, pla, len);
    len = rsa_no_padding((char*)pin, len, 32);
    printf("%d line: %d\n", __LINE__, len);
    int plen = len/sizeof(int);
    for (int i=0; i<plen; i++,len++) {
    #if BAD_WAY
        if (pin[i]/n >255) {
            free(pin);
            free(mc_out);
            return -1;   // please use larger q and p
        }
        mc_out[i*5] = pin[i]/n;
        pout = (int *)(mc_out+i*5+1);
        pout[0] = mod_exp(pin[i], e, n);
    #else 
        if (pin[i] > n) {
            free(pin);
            free(mc_out);
            return -1;
        }
        pout[i] = mod_exp(pin[i], e, n);
    #endif
    }
    // (n+2)*4/3
    base64_encode(mc_out, len, out, &len);

    free(pin);
    free(mc_out);
    return len;
}

int test_decrypt(char *cip, char* out, int len, int d, int n) {
    int *pin = (int *)cip;
    int *pout = (int *)out;
    char *mc_in = (char*)malloc(len);
    printf("%d line: %d\n", __LINE__, len);
    if (base64_decode(cip, len, mc_in, &len)) {
        free(mc_in);
        return -1;
    }
    printf("%d line: %d\n", __LINE__, len);
    for (int i=0; i<len/(sizeof(int)+1); i++) {
    #if BAD_WAY
        pin = (int *)(mc_in+i*5+1);
        pout[i] = mod_exp(pin[0], d, n) + mc_in[i*5]*n;
    #else
        pout[i] = mod_exp(pin[i], d, n);
    #endif
    }
    free(mc_in);
    return rsa_no_unpadding(out, len, 32);
}
*/

/**
 * key format
 * 
 * # common
 * version v0           版本
 * serialNumber 001     序列号
 * hash SHA1            哈希算法 
 * signature xxx        签名
 * name xxx             颁发者名称
 * validity  2024-03-03 有效期
 * 
 * # key info
 * algorithm rsa32b     算法
 * p xxx
 * q xxx
 * n xxx
 * e xxx
 * d xxx
 * padding <no  pkcs1  ...>
 */
#define BEGIN_PRIVATE_HEAD "-----BEGIN RSA32B PRIVATE KEY-----"
#define END_PRIVATE_HEAD "-----END RSA32B PRIVATE KEY-----"
#define BEGIN_PUBLIC_HEAD "-----BEGIN RSA32B PUBLIC KEY-----"
#define END_PUBLIC_HEAD "-----END RSA32B PUBLIC KEY-----"
int rsa32b_key_init_from_certificate(struct rsa32b_key *rkey, char *buff) {
    if (!rkey || !buff) return -1;
    char k[2048];
    int l;
    // printf("buffer: %s\n", buff);
    char *b = strstr(buff, "-----BEGIN");
    if (!b) b = buff;
    else {
        // printf("b1: %s\n", b);
        b = strstr(buff, "\n");
        // printf("b2: %s\n", b);
        b = b+1;
        // printf("b3: %s\n", b);
    }
    char *e = strstr(b, "\n-----END");
    if (!e) {
        e = strstr(b, "\n");
        if (e) l = e-b;
        else l = strlen(b);
    } else {
        l = e-b;
    }
    // printf("b: %s\n, l: %d\n, e: %s\n",b ,l,e);
    if (base64_decode(b, l, k, &l)) return -1;
    k[l] = 0;
    memset(rkey, 0, sizeof(struct rsa32b_key));
    struct config_list cl;
    config_parse_init(&cl);
    config_parse_buffer(&cl, k);
    char *value;
    if (value = config_parse_get_value_by_key(&cl ,"p")) rkey->p = atoi(value); else rkey->p = 0;
    if (value = config_parse_get_value_by_key(&cl ,"q")) rkey->q = atoi(value); else rkey->q = 0;
    if (value = config_parse_get_value_by_key(&cl ,"n")) rkey->n = atoi(value); else rkey->n = 0;
    if (value = config_parse_get_value_by_key(&cl ,"e")) rkey->e = atoi(value); else rkey->e = 0;
    if (value = config_parse_get_value_by_key(&cl ,"d")) rkey->d = atoi(value); else rkey->d = 0;
    if (value = config_parse_get_value_by_key(&cl ,"padding")) rkey->padding = atoi(value); else rkey->padding = 0;
    if (value = config_parse_get_value_by_key(&cl ,"version")) printf("key version: %s\n", value);
    if (value = config_parse_get_value_by_key(&cl ,"name")) printf("name: %s\n", value);
    if (value = config_parse_get_value_by_key(&cl ,"validity")) printf("validity: %s\n", value);
    config_parse_free(&cl);
    printf("===============\n");
    printf("p = %d, q = %d\n", rkey->p, rkey->q);
    printf("n = %d\n", rkey->n);
    printf("e = %d, d = %d\n", rkey->e, rkey->d);
    printf("===============\n");
    return 0;
}

int rsa32b_key_init_from_certificateF(struct rsa32b_key *rkey, char *filename) {
    FILE *fp = fopen(filename, "r");
    if (!fp) {
        perror("fopen");
        return -1;
    }
    char buffer[2048];
    size_t r = fread(buffer, 1, sizeof(buffer), fp);
    buffer[r] = 0;
    return rsa32b_key_init_from_certificate(rkey, buffer);
}

int rsa32b_key_generate_publicKey(struct rsa32b_key *rkey, struct sign_info *info, char *filename) {
    if (!rkey || !filename) return -1;
    if (!rkey->e || !rkey->n) return -2;

    char temp[2048];
    struct config_list rl;
    config_parse_init(&rl);
    config_parse_add_note(&rl, "key info");
    config_parse_add_line(&rl, "algorithm", "rsa32b");
    config_parse_add_line(&rl, "e", itoa(rkey->e, temp, 10));
    config_parse_add_line(&rl, "n", itoa(rkey->n, temp, 10));
    config_parse_add_line(&rl, "padding", int_2_str(rkey->padding, 1));

    if (info->hash) {
        if (config_parse_create_buffer(&rl, temp, sizeof(temp))) {
            config_parse_free(&rl);
            return -3;
        }
        sha1_ctx_t ctx;
        crypto_sha1_init(&ctx);
        crypto_sha1_update(&ctx, temp, strlen(temp));
        crypto_sha1_final(&ctx, temp);

    }

    struct config_list l;
    config_parse_init(&l);
    config_parse_add_note(&l, "common");
    config_parse_add_line(&l, "version", VERSION_RSA_32B);
    if (info->serialNumber) config_parse_add_line(&l, "serialNumber", info->serialNumber);
    if (info->hash) {
        config_parse_add_line(&l, "hash", info->hash);
        config_parse_add_line(&l, "signature", temp);
    }
    if (info->name) config_parse_add_line(&l, "name", info->name);
    if (info->validity) config_parse_add_line(&l, "validity", info->validity);
    config_parse_add_empty(&l);
    config_parse_add_list(&l, &rl);
    if (config_parse_create_buffer(&l, temp, sizeof(temp))) {
        config_parse_free(&l);
        return -3;
    }
    config_parse_free(&l);
    
    char temp2[2048];
    int len = strlen(temp);
    printf("len = %d\n", len);
    if (get_base64_encode_len(len) >= sizeof(temp2)) return -4;
    base64_encode(temp, len, temp2, &len);
    temp2[len] = 0;
    if (len+sizeof(BEGIN_PUBLIC_HEAD)+sizeof(END_PUBLIC_HEAD)+4 >= sizeof(temp)) return -4;

    sprintf(temp, "%s\n%s\n%s", BEGIN_PUBLIC_HEAD, temp2, END_PUBLIC_HEAD);
    FILE *fp = fopen(filename, "w");
    if (!fp) {
        perror("fopen");
        return -5;
    }
    fwrite(temp, 1, strlen(temp), fp);
    fclose(fp);
    return 0;
}

int rsa32b_key_generate_privateKey(struct rsa32b_key *rkey, struct sign_info *info, char *filename) {
    if (!rkey || !filename) return -1;
    if (!rkey->d || !rkey->n) return -2;

    char temp[2048];
    struct config_list rl;
    config_parse_init(&rl);
    config_parse_add_note(&rl, "key info");
    config_parse_add_line(&rl, "algorithm", "rsa32b");
    config_parse_add_line(&rl, "d", itoa(rkey->d, temp, 10));
    config_parse_add_line(&rl, "n", itoa(rkey->n, temp, 10));
    config_parse_add_line(&rl, "padding", int_2_str(rkey->padding, 1));

    if (info->hash) {
        if (config_parse_create_buffer(&rl, temp, sizeof(temp))) {
            config_parse_free(&rl);
            return -3;
        }
        sha1_ctx_t ctx;
        crypto_sha1_init(&ctx);
        crypto_sha1_update(&ctx, temp, strlen(temp));
        crypto_sha1_final(&ctx, temp);

    }
    
    struct config_list l;
    config_parse_init(&l);
    config_parse_add_note(&l, "common");
    config_parse_add_line(&l, "version", VERSION_RSA_32B);
    if (info->serialNumber) config_parse_add_line(&l, "serialNumber", info->serialNumber);
    if (info->hash) {
        config_parse_add_line(&l, "hash", info->hash);
        config_parse_add_line(&l, "signature", temp);
    }
    if (info->name) config_parse_add_line(&l, "name", info->name);
    if (info->validity) config_parse_add_line(&l, "validity", info->validity);
    config_parse_add_empty(&l);
    config_parse_add_list(&l, &rl);
    if (config_parse_create_buffer(&l, temp, sizeof(temp))) {
        config_parse_free(&l);
        return -3;
    }
    config_parse_free(&l);

    char temp2[2048];
    int len = strlen(temp);
    if (get_base64_encode_len(len) >= sizeof(temp2)) return -4;
    base64_encode(temp, len, temp2, &len);
    temp2[len] = 0;
    if (len+sizeof(BEGIN_PRIVATE_HEAD)+sizeof(END_PRIVATE_HEAD)+4 >= sizeof(temp)) return -4;

    sprintf(temp, "%s\n%s\n%s", BEGIN_PRIVATE_HEAD, temp2, END_PRIVATE_HEAD);
    FILE *fp = fopen(filename, "w");
    if (!fp) {
        perror("fopen");
        return -5;
    }
    fwrite(temp, 1, strlen(temp), fp);
    fclose(fp);

    return 0;
}

int rsa32b_key_generate_ca(struct rsa32b_key *rkey, struct sign_info *info, char *filename) {
    if (!rkey || !filename) return -1;
    if (!rkey->d || !rkey->n) return -2;

    char temp[2048];
    struct config_list rl;
    config_parse_init(&rl);
    config_parse_add_note(&rl, "key info");
    config_parse_add_line(&rl, "algorithm", "rsa32b");
    config_parse_add_line(&rl, "p", itoa(rkey->p, temp, 10));
    config_parse_add_line(&rl, "q", itoa(rkey->q, temp, 10));
    config_parse_add_line(&rl, "e", itoa(rkey->e, temp, 10));
    config_parse_add_line(&rl, "d", itoa(rkey->d, temp, 10));
    config_parse_add_line(&rl, "n", itoa(rkey->n, temp, 10));
    config_parse_add_line(&rl, "padding", int_2_str(rkey->padding, 1));

    if (info->hash) {
        if (config_parse_create_buffer(&rl, temp, sizeof(temp))) {
            config_parse_free(&rl);
            return -3;
        }
        sha1_ctx_t ctx;
        crypto_sha1_init(&ctx);
        crypto_sha1_update(&ctx, temp, strlen(temp));
        crypto_sha1_final(&ctx, temp);

    }
    
    struct config_list l;
    config_parse_init(&l);
    config_parse_add_note(&l, "common");
    config_parse_add_line(&l, "version", VERSION_RSA_32B);
    if (info->serialNumber) config_parse_add_line(&l, "serialNumber", info->serialNumber);
    if (info->hash) {
        config_parse_add_line(&l, "hash", info->hash);
        config_parse_add_line(&l, "signature", temp);
    }
    if (info->name) config_parse_add_line(&l, "name", info->name);
    if (info->validity) config_parse_add_line(&l, "validity", info->validity);
    config_parse_add_empty(&l);
    config_parse_add_list(&l, &rl);
    if (config_parse_create_buffer(&l, temp, sizeof(temp))) {
        config_parse_free(&l);
        return -3;
    }
    config_parse_free(&l);

    char temp2[2048];
    int len = strlen(temp);
    if (get_base64_encode_len(len) >= sizeof(temp2)) return -4;
    base64_encode(temp, len, temp2, &len);
    temp2[len] = 0;
    if (len+sizeof(BEGIN_PRIVATE_HEAD)+sizeof(END_PRIVATE_HEAD)+4 >= sizeof(temp)) return -4;

    sprintf(temp, "%s\n%s\n%s", BEGIN_PRIVATE_HEAD, temp2, END_PRIVATE_HEAD);
    FILE *fp = fopen(filename, "w");
    if (!fp) {
        perror("fopen");
        return -5;
    }
    fwrite(temp, 1, strlen(temp), fp);
    fclose(fp);

    return 0;
}

int main() {
    int n=264539167, e=65537, d=143050593;

    printf("Public Key: (n = %d, e = %d)\n", n, e);
    printf("Private Key: (n = %d, d = %d)\n", n, d);

    int plaintext = 264539169;  // 明文字符 'A'
    printf("Plaintext: %d\n", plaintext);

    int ciphertext = rsa_encrypt_simple(plaintext, e, n);
    printf("Ciphertext: %d\n", ciphertext);

    int decryptedtext = rsa_decrypt_simple(ciphertext, d, n);
    printf("Decrypted Text: %d\n", decryptedtext);

    struct rsa32b_key rkey;
    int ret;
    if (ret = rsa32b_key_create(&rkey, 8783, 22613, 0)) {
        printf("create key fail: %d\n", ret);
        return -1;
    }
    char *arr_pla = "qwerasdfzxcv vcxzfdaareqe 不擦检查完 9hd()(&()&(cadamcw))";
    char cip[200] = {0};
    char pla[200] = {0};
    ret = rsa32b_private_encrypt(&rkey, arr_pla, cip, strlen(arr_pla));
    if (ret <= 0) {
        printf("rsa32b_private_encrypt fail: %d\n", ret);
        return -1;
    }
    printf("cip(%d): %s\n", ret, cip);
    ret = rsa32b_public_decrypt(&rkey, cip, pla, ret);
    if (ret <= 0) {
        printf("rsa32b_public_decrypt fail: %d\n", ret);
        return -1;
    }
    printf("pla(%d): %s\n", ret, pla);

    // struct sign_info info = {"123", "SHA1", "unknow", "2023-04-12"};
    // if (ret = rsa32b_key_generate_publicKey(&rkey, &info, "public.key")) {
    //     printf("rsa32b_key_generate_publicKey fail: %d\n", ret);
    // }
    // if (ret = rsa32b_key_generate_privateKey(&rkey, &info, "private.key")) {
    //     printf("rsa32b_key_generate_privateKey fail: %d\n", ret);
    // }
    // rsa32b_key_generate_ca(&rkey, &info, "ca.key");

    struct rsa32b_key r2key;
    if (ret = rsa32b_key_init_from_certificateF(&r2key, "public.key")) {
        printf("rsa32b_key_init_from_certificateF fail: %d\n", ret);
    }
    if ((ret = rsa32b_public_encrypt(&r2key, "123456", cip, 6))<=0) {
        printf("rsa32b_public_encrypt fail: %d\n", ret);
    }
    printf("123456 --> %s\n", cip);
    if (ret = rsa32b_key_init_from_certificateF(&r2key, "private.key")) {
        printf("rsa32b_key_init_from_certificateF fail: %d\n", ret);
    }
    if ((ret = rsa32b_private_decrypt(&r2key, cip, pla, strlen(cip)))<=0) {
        printf("rsa32b_private_decrypt fail: %d\n", ret);
    }
    printf("%s --> %s\n", cip, pla);

    

    return 0;
}
